import tushare as ts
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from datetime import datetime, timedelta
import warnings

warnings.filterwarnings('ignore')

# 初始化pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')


class StockAnalyzer:
    def __init__(self, ts_code, start_date, end_date, freq='D'):
        """
        初始化股票分析器
        :param ts_code: 股票代码，格式 '000001.SZ'
        :param start_date: 开始日期，格式 '20180101'
        :param end_date: 结束日期，格式 '20201231'
        :param freq: 时间频率，'D'日线，'W'周线，'M'月线，'60'60分钟线等
        """
        self.ts_code = ts_code
        self.start_date = start_date
        self.end_date = end_date
        self.freq = freq
        self.data = None
        self.backtest_result = None
        self.name = None
        self.trade_history = None
        self.trade_stats = None

    def get_data(self):
        """获取股票数据"""
        if self.freq in ['D', 'W', 'M']:  # 日线、周线、月线
            df = pro.daily(ts_code=self.ts_code,
                           start_date=self.start_date,
                           end_date=self.end_date)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            df.set_index('trade_date', inplace=True)

            if self.freq == 'W':
                df = df.resample('W').last()
            elif self.freq == 'M':
                df = df.resample('M').last()

        else:  # 分钟线
            df = pro.bbar(ts_code=self.ts_code,
                          freq=self.freq,
                          start_date=self.start_date,
                          end_date=self.end_date)
            df['trade_time'] = pd.to_datetime(df['trade_time'])
            df.set_index('trade_time', inplace=True)

        df = df.sort_index()
        self.data = df
        self.name = pro.stock_basic(ts_code=self.ts_code)['name'].values[0]
        return df

    def calculate_technical_indicators(self):
        """计算技术指标"""
        if self.data is None:
            self.get_data()

        df = self.data.copy()

        # 计算MACD
        df['EMA快线'] = df['close'].ewm(span=12, adjust=False).mean()
        df['EMA慢线'] = df['close'].ewm(span=26, adjust=False).mean()
        df['DIF'] = df['EMA快线'] - df['EMA慢线']
        df['DEA'] = df['DIF'].ewm(span=9, adjust=False).mean()
        df['MACD'] = (df['DIF'] - df['DEA']) * 2

        # 计算RSI
        delta = df['close'].diff()
        gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
        rs = gain / loss
        df['RSI'] = 100 - (100 / (1 + rs))

        # 计算布林带
        df['MA20'] = df['close'].rolling(window=20).mean()
        df['Upper'] = df['MA20'] + 2 * df['close'].rolling(window=20).std()
        df['Lower'] = df['MA20'] - 2 * df['close'].rolling(window=20).std()

        self.data = df
        return df

    def generate_signals(self, stop_loss_pct=0.05, stop_gain_pct=0.15):
        """生成交易信号，包含止损止盈"""
        if self.data is None:
            self.calculate_technical_indicators()

        df = self.data.copy()
        df['信号'] = 0
        df['止损价'] = np.nan
        df['止盈价'] = np.nan

        position = 0  # 0表示空仓，1表示持仓
        stop_loss_price = 0
        take_profit_price = 0

        for i in range(1, len(df)):
            # MACD金叉买入信号
            if position == 0 and df['DIF'].iloc[i] > df['DEA'].iloc[i] and df['DIF'].iloc[i - 1] <= df['DEA'].iloc[
                i - 1]:
                df.iloc[i, df.columns.get_loc('信号')] = 1
                position = 1
                stop_loss_price = df['close'].iloc[i] * (1 - stop_loss_pct)
                take_profit_price = df['close'].iloc[i] * (1 + stop_gain_pct)
                df.iloc[i, df.columns.get_loc('止损价')] = stop_loss_price
                df.iloc[i, df.columns.get_loc('止盈价')] = take_profit_price

            # 止损或止盈卖出信号
            elif position == 1:
                if df['low'].iloc[i] <= stop_loss_price:  # 触发止损
                    df.iloc[i, df.columns.get_loc('信号')] = -2  # -2表示止损卖出
                    position = 0
                elif df['high'].iloc[i] >= take_profit_price:  # 触发止盈
                    df.iloc[i, df.columns.get_loc('信号')] = -3  # -3表示止盈卖出
                    position = 0
                elif df['DIF'].iloc[i] < df['DEA'].iloc[i] and df['DIF'].iloc[i - 1] >= df['DEA'].iloc[i - 1]:  # MACD死叉
                    df.iloc[i, df.columns.get_loc('信号')] = -1  # -1表示普通卖出
                    position = 0

        self.data = df
        return df

    def backtest(self, initial_capital=100000, commission=0.0003):
        """回测策略"""
        if self.data is None:
            self.generate_signals()

        df = self.data.copy()
        capital = initial_capital
        shares = 0
        position = 0  # 0表示空仓，1表示持仓
        trade_history = []

        backtest_df = pd.DataFrame(index=df.index)
        backtest_df['收盘价'] = df['close']
        backtest_df['信号'] = df['信号']
        backtest_df['现金'] = capital
        backtest_df['持仓'] = shares
        backtest_df['总资产'] = capital

        for i in range(len(df)):
            current_date = df.index[i]
            current_price = df['close'].iloc[i]
            signal = df['信号'].iloc[i]

            # 买入信号
            if signal == 1 and position == 0:
                # 计算可买数量（以100股为单位）
                max_shares = int((capital * (1 - commission)) / (current_price * 100)) * 100
                if max_shares > 0:
                    shares = max_shares
                    cost = shares * current_price * (1 + commission)
                    capital -= cost
                    position = 1
                    trade_history.append({
                        '日期': current_date,
                        '类型': '买入',
                        '价格': round(current_price, 2),
                        '数量': shares,
                        '金额': round(-cost, 2),
                        '现金余额': round(capital, 2),
                        '持仓数量': shares,
                        '信号类型': 'MACD金叉'
                    })

            # 卖出信号（普通卖出、止损、止盈）
            elif signal in [-1, -2, -3] and position == 1:
                proceeds = shares * current_price * (1 - commission)
                capital += proceeds

                trade_type = '卖出'
                signal_type = 'MACD死叉'
                if signal == -2:
                    trade_type = '止损卖出'
                    signal_type = '止损触发'
                elif signal == -3:
                    trade_type = '止盈卖出'
                    signal_type = '止盈触发'

                trade_history.append({
                    '日期': current_date,
                    '类型': trade_type,
                    '价格': round(current_price, 2),
                    '数量': shares,
                    '金额': round(proceeds, 2),
                    '现金余额': round(capital, 2),
                    '持仓数量': 0,
                    '信号类型': signal_type
                })
                shares = 0
                position = 0

            # 更新每日资产
            if position == 1:
                total = capital + shares * current_price
            else:
                total = capital

            backtest_df.iloc[i, backtest_df.columns.get_loc('现金')] = capital
            backtest_df.iloc[i, backtest_df.columns.get_loc('持仓')] = shares
            backtest_df.iloc[i, backtest_df.columns.get_loc('总资产')] = total

        self.backtest_result = backtest_df
        self.trade_history = pd.DataFrame(trade_history)
        return backtest_df

    def calculate_performance(self):
        """计算策略表现"""
        if self.backtest_result is None:
            self.backtest()

        df = self.backtest_result
        initial_capital = df['现金'].iloc[0]
        final_asset = df['总资产'].iloc[-1]
        strategy_return_pct = (final_asset - initial_capital) / initial_capital * 100

        # 计算年化收益率
        days = (df.index[-1] - df.index[0]).days
        annualized_return = (final_asset / initial_capital) ** (365 / days) - 1 if days > 0 else 0

        # 计算买入持有策略收益
        buy_hold_return_pct = (df['收盘价'].iloc[-1] - df['收盘价'].iloc[0]) / df['收盘价'].iloc[0] * 100

        # 计算最大回撤
        df['峰值'] = df['总资产'].cummax()
        df['回撤'] = (df['峰值'] - df['总资产']) / df['峰值']
        max_drawdown_pct = df['回撤'].max() * 100

        # 计算交易次数
        trade_count = len(self.trade_history) // 2 if len(self.trade_history) > 0 else 0  # 每次完整买卖算一次交易

        # 计算胜率和平均盈亏比
        if trade_count > 0:
            winning_trades = 0
            total_profit = 0
            total_loss = 0
            profit_count = 0
            loss_count = 0

            for i in range(0, len(self.trade_history) - 1, 2):
                buy_trade = self.trade_history.iloc[i]
                sell_trade = self.trade_history.iloc[i + 1]

                profit = sell_trade['金额'] + buy_trade['金额']
                if profit > 0:
                    winning_trades += 1
                    total_profit += profit
                    profit_count += 1
                else:
                    total_loss += abs(profit)
                    loss_count += 1

            win_rate = winning_trades / trade_count * 100 if trade_count > 0 else 0
            avg_profit = total_profit / profit_count if profit_count > 0 else 0
            avg_loss = total_loss / loss_count if loss_count > 0 else 0
            profit_loss_ratio = avg_profit / avg_loss if avg_loss > 0 else 0
        else:
            win_rate = 0
            profit_loss_ratio = 0

        # 添加交易统计
        if not self.trade_history.empty:
            # 计算每笔交易持有天数
            trades = self.trade_history.copy()
            trades['交易ID'] = (trades['类型'] == '买入').cumsum()
            trade_stats = trades.groupby('交易ID').agg({
                '日期': ['first', 'last'],
                '金额': ['first', 'last'],
                '类型': 'last'
            })

            trade_stats.columns = ['买入日期', '卖出日期', '买入金额', '卖出金额', '卖出类型']
            trade_stats['持有天数'] = (trade_stats['卖出日期'] - trade_stats['买入日期']).dt.days
            trade_stats['盈亏'] = trade_stats['卖出金额'] + trade_stats['买入金额']
            trade_stats['收益率'] = trade_stats['盈亏'] / (-trade_stats['买入金额'])

            self.trade_stats = trade_stats
        else:
            self.trade_stats = pd.DataFrame()

        performance = {
            '初始资金': initial_capital,
            '最终资产': final_asset,
            '策略收益率(%)': strategy_return_pct,
            '年化收益率(%)': annualized_return * 100,
            '买入持有收益率(%)': buy_hold_return_pct,
            '最大回撤(%)': max_drawdown_pct,
            '交易次数': trade_count,
            '胜率(%)': win_rate,
            '平均盈亏比': profit_loss_ratio
        }

        return performance

    def plot_results(self):
        """绘制结果图表"""
        if self.backtest_result is None:
            self.backtest()

        df = self.data
        backtest_df = self.backtest_result
        performance = self.calculate_performance()

        # 创建子图
        fig = make_subplots(rows=4, cols=1,
                            shared_xaxes=True,
                            vertical_spacing=0.05,
                            subplot_titles=(
                                f'{self.name} ({self.ts_code}) 价格与交易信号',
                                'MACD指标',
                                '布林带指标',
                                '资金曲线'),
                            row_heights=[0.4, 0.2, 0.2, 0.2])

        # 价格和信号图
        fig.add_trace(go.Scatter(x=df.index, y=df['close'],
                                 name='收盘价', line=dict(color='blue')),
                      row=1, col=1)

        # 买入信号
        buy_signals = df[df['信号'] == 1]
        fig.add_trace(go.Scatter(x=buy_signals.index, y=buy_signals['close'],
                                 mode='markers', name='买入',
                                 marker=dict(color='green', size=10, symbol='triangle-up')),
                      row=1, col=1)

        # 卖出信号
        sell_signals = df[df['信号'] == -1]
        fig.add_trace(go.Scatter(x=sell_signals.index, y=sell_signals['close'],
                                 mode='markers', name='普通卖出',
                                 marker=dict(color='red', size=10, symbol='triangle-down')),
                      row=1, col=1)

        # 止损信号
        stop_loss_signals = df[df['信号'] == -2]
        if not stop_loss_signals.empty:
            fig.add_trace(go.Scatter(x=stop_loss_signals.index, y=stop_loss_signals['close'],
                                     mode='markers', name='止损卖出',
                                     marker=dict(color='darkred', size=10, symbol='x')),
                          row=1, col=1)

        # 止盈信号
        take_profit_signals = df[df['信号'] == -3]
        if not take_profit_signals.empty:
            fig.add_trace(go.Scatter(x=take_profit_signals.index, y=take_profit_signals['close'],
                                     mode='markers', name='止盈卖出',
                                     marker=dict(color='purple', size=10, symbol='star')),
                          row=1, col=1)

        # 布林带
        fig.add_trace(go.Scatter(x=df.index, y=df['Upper'],
                                 name='上轨', line=dict(color='gray', width=1)),
                      row=3, col=1)
        fig.add_trace(go.Scatter(x=df.index, y=df['Lower'],
                                 name='下轨', line=dict(color='gray', width=1)),
                      row=3, col=1)
        fig.add_trace(go.Scatter(x=df.index, y=df['MA20'],
                                 name='中轨', line=dict(color='blue', width=1)),
                      row=3, col=1)
        fig.add_trace(go.Scatter(x=df.index, y=df['close'],
                                 name='收盘价', line=dict(color='orange', width=1)),
                      row=3, col=1)

        # MACD图
        fig.add_trace(go.Scatter(x=df.index, y=df['DIF'],
                                 name='DIF', line=dict(color='blue')),
                      row=2, col=1)
        fig.add_trace(go.Scatter(x=df.index, y=df['DEA'],
                                 name='DEA', line=dict(color='orange')),
                      row=2, col=1)

        # MACD柱状图
        colors = np.where(df['MACD'] > 0, 'red', 'green')
        fig.add_trace(go.Bar(x=df.index, y=df['MACD'],
                             name='MACD', marker_color=colors),
                      row=2, col=1)

        # 资金曲线图
        fig.add_trace(go.Scatter(x=backtest_df.index, y=backtest_df['总资产'],
                                 name='总资产', line=dict(color='purple')),
                      row=4, col=1)

        # 添加基准线（初始资金）
        fig.add_trace(go.Scatter(x=backtest_df.index,
                                 y=[performance['初始资金']] * len(backtest_df),
                                 name='初始资金',
                                 line=dict(color='gray', dash='dash')),
                      row=4, col=1)

        # 更新布局
        fig.update_layout(
            height=1200,
            title_text=f"股票交易策略分析 - {self.name} ({self.ts_code}) {self.start_date}至{self.end_date} {self.freq}线",
            hovermode="x unified",
            showlegend=True
        )

        # 更新y轴标题
        fig.update_yaxes(title_text="价格(元)", row=1, col=1)
        fig.update_yaxes(title_text="MACD", row=2, col=1)
        fig.update_yaxes(title_text="布林带", row=3, col=1)
        fig.update_yaxes(title_text="资产(元)", row=4, col=1)

        # 添加策略表现注释
        performance_text = (
            f"初始资金: {performance['初始资金']:,.2f}元 | "
            f"最终资产: {performance['最终资产']:,.2f}元<br>"
            f"策略收益率: {performance['策略收益率(%)']:.2f}% | "
            f"年化收益率: {performance['年化收益率(%)']:.2f}%<br>"
            f"买入持有收益率: {performance['买入持有收益率(%)']:.2f}% | "
            f"最大回撤: {performance['最大回撤(%)']:.2f}%<br>"
            f"交易次数: {performance['交易次数']} | "
            f"胜率: {performance['胜率(%)']:.2f}% | "
            f"平均盈亏比: {performance['平均盈亏比']:.2f}"
        )

        fig.add_annotation(
            xref="paper", yref="paper",
            x=0.02, y=1.08,
            text=performance_text,
            showarrow=False,
            font=dict(size=12),
            align="left"
        )

        return fig

    def show_trade_history(self):
        """显示交易历史记录"""
        if self.trade_history is None:
            self.backtest()

        pd.set_option('display.max_columns', None)
        pd.set_option('display.width', 1000)
        pd.set_option('display.float_format', '{:.2f}'.format)

        if self.trade_history.empty:
            print("\n没有交易记录")
            return None

        print("\n完整交易历史记录:")
        print(self.trade_history.to_string(index=False))

        if hasattr(self, 'trade_stats') and not self.trade_stats.empty:
            print("\n交易统计:")
            print(self.trade_stats.to_string(index=False))

        return self.trade_history


def compare_strategies(stock_codes, start_date, end_date, freq='D'):
    """
    比较多个股票的策略表现
    :param stock_codes: 股票代码列表，如 ['000001.SZ', '600000.SH']
    :param start_date: 开始日期
    :param end_date: 结束日期
    :param freq: 时间频率
    """
    analyzers = []
    performances = []

    for code in stock_codes:
        analyzer = StockAnalyzer(code, start_date, end_date, freq)
        analyzer.get_data()
        analyzer.calculate_technical_indicators()
        analyzer.generate_signals()
        analyzer.backtest()
        performance = analyzer.calculate_performance()
        performance['代码'] = code
        performance['名称'] = analyzer.name
        performances.append(performance)
        analyzers.append(analyzer)

    # 创建比较图表
    fig = go.Figure()

    for analyzer in analyzers:
        fig.add_trace(go.Scatter(
            x=analyzer.backtest_result.index,
            y=analyzer.backtest_result['总资产'],
            name=f"{analyzer.name} ({analyzer.ts_code})",
            mode='lines'
        ))

    # 添加基准线（初始资金）
    fig.add_trace(go.Scatter(
        x=analyzers[0].backtest_result.index,
        y=[analyzers[0].backtest_result['总资产'].iloc[0]] * len(analyzers[0].backtest_result),
        name='初始资金',
        line=dict(color='gray', dash='dash')
    ))

    # 准备性能比较表格
    perf_df = pd.DataFrame(performances)
    perf_df = perf_df[['名称', '代码', '初始资金', '最终资产', '策略收益率(%)',
                       '年化收益率(%)', '买入持有收益率(%)', '最大回撤(%)',
                       '交易次数', '胜率(%)', '平均盈亏比']]

    # 格式化数字显示
    perf_df['初始资金'] = perf_df['初始资金'].apply(lambda x: f"{x:,.2f}")
    perf_df['最终资产'] = perf_df['最终资产'].apply(lambda x: f"{x:,.2f}")
    perf_df['策略收益率(%)'] = perf_df['策略收益率(%)'].apply(lambda x: f"{x:.2f}%")
    perf_df['年化收益率(%)'] = perf_df['年化收益率(%)'].apply(lambda x: f"{x:.2f}%")
    perf_df['买入持有收益率(%)'] = perf_df['买入持有收益率(%)'].apply(lambda x: f"{x:.2f}%")
    perf_df['最大回撤(%)'] = perf_df['最大回撤(%)'].apply(lambda x: f"{x:.2f}%")
    perf_df['胜率(%)'] = perf_df['胜率(%)'].apply(lambda x: f"{x:.2f}%")
    perf_df['平均盈亏比'] = perf_df['平均盈亏比'].apply(lambda x: f"{x:.2f}")

    fig.update_layout(
        title_text=f"多股票策略比较 {start_date}至{end_date} {freq}线",
        xaxis_title="日期",
        yaxis_title="资产价值(元)",
        hovermode="x unified",
        height=600
    )

    return fig, perf_df


# 示例使用
if __name__ == "__main__":
    # 设置时间范围
    start_date = '20200101'
    end_date = '20231231'

    print("=== 单股票分析示例 ===")
    # 示例1: 单股票分析
    analyzer = StockAnalyzer('600519.SH', start_date, end_date, 'D')  # 贵州茅台日线
    analyzer.get_data()
    analyzer.calculate_technical_indicators()
    analyzer.generate_signals(stop_loss_pct=0.07, stop_gain_pct=0.2)  # 7%止损，20%止盈
    analyzer.backtest()
    performance = analyzer.calculate_performance()
    fig = analyzer.plot_results()
    fig.show()

    # 显示交易历史
    analyzer.show_trade_history()

    # 打印性能指标
    print("\n策略表现:")
    for k, v in performance.items():
        if '收益率' in k or '回撤' in k or '胜率' in k:
            print(f"{k}: {v:.2f}%")
        elif '初始资金' in k or '最终资产' in k:
            print(f"{k}: {v:,.2f}元")
        else:
            print(f"{k}: {v}")

    print("\n=== 多股票比较示例 ===")
    # 示例2: 多股票比较
    stock_codes = ['000001.SZ', '600000.SH', '601318.SH']  # 平安银行、浦发银行、中国平安
    compare_fig, perf_table = compare_strategies(stock_codes, start_date, end_date, 'W')  # 周线比较
    compare_fig.show()

    print("\n多股票策略比较:")
    print(perf_table.to_string(index=False))